#pragma once

#include <Eigen/Core>
#include <unsupported/Eigen/MatrixFunctions>

#include "stair_ctrl/interface.hpp"
#include "stair_ctrl/math.hpp"

// state estimator parameters
#define STATE_SIZE 18
#define MEAS_SIZE 24
#define NUM_LEG 4

// implement a basic error state KF to estimate robot pose
// assume orientation is known from a IMU (state.root_rot_mat)
template <typename T>
class extendedKalmanFilter
{
public:
    extendedKalmanFilter()
    {
        // constructor
        eye3.setIdentity();
        // C is fixed
        C.setZero();
        for (int i = 0; i < NUM_LEG; ++i)
        {
            C.template block<3, 3>(i * 3, 0) = eye3;              //pos
            C.template block<3, 3>(i * 3, 6 + i * 3) = -eye3;       //-foot pos
            C.template block<3, 3>(NUM_LEG * 3 + i * 3, 3) = eye3; // vel
            // C(NUM_LEG * 6 + i, 6 + i * 3 + 2) = 1;        // height z of foot
        }

        // Q R are fixed
        Q.setIdentity();
        Q.template block<3, 3>(0, 0) = process_noise_p_imu * eye3; // position transition
        Q.template block<3, 3>(3, 3) = process_noise_v_imu * eye3; // velocity transition
        for (int i = 0; i < NUM_LEG; ++i)
        {
            Q.template block<3, 3>(6 + i * 3, 6 + i * 3) = process_noise_p_foot * eye3; // foot position transition
        }

        R.setIdentity();
        for (int i = 0; i < NUM_LEG; ++i)
        {
            R.template block<3, 3>(i * 3, i * 3) = sensor_noise_p_imu_rel_foot * eye3;                             // fk estimation
            R.template block<3, 3>(NUM_LEG * 3 + i * 3, NUM_LEG * 3 + i * 3) = sensor_noise_v_imu_rel_foot * eye3; // vel estimation
            // R(NUM_LEG * 6 + i, NUM_LEG * 6 + i) = sensor_noise_z_foot;                                    // height z estimation
        }

        // set A to identity
        A.setIdentity();

        // set B to zero
        B.setZero();

        // assume_flat_ground = false;

        // if (assume_flat_ground == false)
        // {
        //     for (int i = 0; i < NUM_LEG; ++i)
        //     {
        //         R(NUM_LEG * 6 + i, NUM_LEG * 6 + i) = 1e5; // height z estimation not reliable
        //     }
        // }
    }

    void setParam(T process_noise_p_imu_,
                    T process_noise_v_imu_,
                    T process_noise_p_foot_,
                    T sensor_noise_p_imu_rel_foot_,
                    T sensor_noise_v_imu_rel_foot_,
                    T sensor_noise_z_foot_,
                    T high_suspect_num_)
    {
        process_noise_p_imu = process_noise_p_imu_;
        process_noise_v_imu = process_noise_v_imu_;
        process_noise_p_foot = process_noise_p_foot_;
        sensor_noise_p_imu_rel_foot = sensor_noise_p_imu_rel_foot_;
        sensor_noise_v_imu_rel_foot = sensor_noise_v_imu_rel_foot_;
        sensor_noise_z_foot = sensor_noise_z_foot_;
        high_suspect_num = high_suspect_num_;
        
        // refresh Q and R
        Q.setIdentity();
        Q.template block<3, 3>(0, 0) = process_noise_p_imu * eye3; // position transition
        Q.template block<3, 3>(3, 3) = process_noise_v_imu * eye3; // velocity transition
        for (int i = 0; i < NUM_LEG; ++i)
        {
            Q.template block<3, 3>(6 + i * 3, 6 + i * 3) = process_noise_p_foot * eye3; // foot position transition
        }

        R.setIdentity();
        for (int i = 0; i < NUM_LEG; ++i)
        {
            R.template block<3, 3>(i * 3, i * 3) = sensor_noise_p_imu_rel_foot * eye3;                             // fk estimation
            R.template block<3, 3>(NUM_LEG * 3 + i * 3, NUM_LEG * 3 + i * 3) = sensor_noise_v_imu_rel_foot * eye3; // vel estimation
            // R(NUM_LEG * 6 + i, NUM_LEG * 6 + i) = sensor_noise_z_foot;                                    // height z estimation
        }

        // if (assume_flat_ground == false)
        // {
        //     for (int i = 0; i < NUM_LEG; ++i)
        //     {
        //         R(NUM_LEG * 6 + i, NUM_LEG * 6 + i) = 1e5; // height z estimation not reliable
        //     }
        // }
        // std::cout << "Q" << std::endl
        //             << Q << std::endl
        //             << "R" << std::endl
        //             << R << std::endl;
    }

    void init_state(robotInterface<T> &state)
    {
        filter_initialized = true;
        P.setIdentity();
        P = P * 100;

        // set initial value of x
        x.setZero();
        x.template segment<3>(0) = Eigen::Matrix<T, 3, 1>(0, 0, 0.09);
        auto foot_pose_rel = state.getFootPosBody();
        auto rotmat = state.getRotMat(state.body_quat).transpose();
        for (int i = 0; i < NUM_LEG; ++i)
        {
            Eigen::Matrix<T, 3, 1> fk_pos = foot_pose_rel.template segment<3>(3 * i);
            x.template segment<3>(6 + i * 3) = rotmat * fk_pos + x.template segment<3>(0);
        }
    }

    void update_estimation(robotInterface<T> &state, T dt)
    {
        // update A B using latest dt
        A.template block<3, 3>(0, 3) = dt * eye3;
        B.template block<3, 3>(0, 0) = 0.5 * dt * dt * eye3;
        B.template block<3, 3>(3, 0) = dt * eye3;

        // control input u is Ra + ag
        auto rotmat = state.getRotMat(state.body_quat).transpose();
        Eigen::Matrix<T, 3, 1> u = rotmat * state.body_lin_acc + Eigen::Matrix<T, 3, 1>(0, 0, -gravity);

        // contact estimation, do something very simple first
        // if (state.movement_mode == 0) {  // stand
        //     for (int i = 0; i < NUM_LEG; ++i) estimated_contacts[i] = 1.0;
        // } else {  // walk
        //     for (int i = 0; i < NUM_LEG; ++i) {
        //         estimated_contacts[i] = std::min(std::max((state.foot_force(i)) / (100.0 - 0.0), 0.0), 1.0);
        //         // estimated_contacts[i] = 1.0/(1.0+std::exp(-(state.foot_force(i)-100)));
        //     }
        // }
        for (int i = 0; i < NUM_LEG; ++i)
        {
            estimated_contacts[i] = state.contact_status[i] ? 1.0 : 0.0;
        }

        // update Q
        Q.template block<3, 3>(0, 0) = process_noise_p_imu * eye3;
        Q.template block<3, 3>(3, 3) = process_noise_v_imu * eye3;
        // update Q R for legs not in contact
        for (int i = 0; i < NUM_LEG; ++i)
        {
            Q.template block<3, 3>(6 + i * 3, 6 + i * 3) =
                (1 + (1 - estimated_contacts[i]) * high_suspect_num) * process_noise_p_foot * eye3; // foot position transition
            // for estimated_contacts[i] == 1, Q = 0.002
            // for estimated_contacts[i] == 0, Q = 1001*Q

            R.template block<3, 3>(i * 3, i * 3) = (1 + (1 - estimated_contacts[i]) * high_suspect_num) * sensor_noise_p_imu_rel_foot * eye3; // fk estimation
            // R.template block<3, 3>(i * 3, i * 3) = sensor_noise_p_imu_rel_foot * eye3; // fk estimation
            R.template block<3, 3>(NUM_LEG * 3 + i * 3, NUM_LEG * 3 + i * 3) = (1 + (1 - estimated_contacts[i]) * high_suspect_num) * sensor_noise_v_imu_rel_foot * eye3; // vel estimation
            // if (assume_flat_ground)
            // {
            //     R(NUM_LEG * 6 + i, NUM_LEG * 6 + i) = (1 + (1 - estimated_contacts[i]) * high_suspect_num) * sensor_noise_z_foot; // height z estimation
            // }
        }

        // process update
        xbar = A * x + B * u;
        Pbar = A * P * A.transpose() + Q;

        // measurement construction
        yhat = C * xbar;
        // leg_v = (-J_rf*av-skew(omega)*p_rf);
        // r((i-1)*3+1:(i-1)*3+3) = body_v - R_er*leg_v;
        // actual measurement
        auto foot_pos_rel = state.getFootPosBody();
        auto foot_vel_rel = state.getFootVelBody();

        // for (int i = 0; i < NUM_LEG; ++i)
        // {
        //     Eigen::Matrix<T, 3, 1> fk_pos = foot_pos_rel.template segment<3>(3 * i);
        //     y.template block<3, 1>(i * 3, 0) = rotmat * fk_pos; // fk estimation
        //     Eigen::Matrix<T, 3, 1> leg_v = -foot_vel_rel.template segment<3>(3 * i) - skew(state.body_ang_vel) * fk_pos;
        //     // Eigen::Matrix<T, 3, 1> leg_v = foot_vel_rel.template segment<3>(3 * i) + skew(state.body_ang_vel) * fk_pos;
        //     y.template block<3, 1>(NUM_LEG * 3 + i * 3, 0) = rotmat.transpose() * fk_pos;

        //     y(NUM_LEG * 6 + i) = (1.0 - estimated_contacts[i]) * (x(2) + fk_pos(2)) + estimated_contacts[i] * 0; // height z estimation
        // }
        // std::cout << "y: " << y.head(12).transpose() << std::endl
        //             << "R^T * f: " << y.segment(12, 12).transpose() << std::endl;
        for (int i = 0; i < NUM_LEG; ++i)
        {
            Eigen::Matrix<T, 3, 1> fk_pos = foot_pos_rel.template segment<3>(3 * i);
            y.template block<3, 1>(i * 3, 0) = -rotmat * fk_pos; // fk estimation
            Eigen::Matrix<T, 3, 1> leg_v = -foot_vel_rel.template segment<3>(3 * i) - skew(state.body_ang_vel) * fk_pos;
            // Eigen::Matrix<T, 3, 1> leg_v = foot_vel_rel.template segment<3>(3 * i) + skew(state.body_ang_vel) * fk_pos;
            y.template block<3, 1>(NUM_LEG * 3 + i * 3, 0) = (1.0 - estimated_contacts[i]) * x.template segment<3>(3) + estimated_contacts[i] * rotmat * leg_v; // vel estimation

            // y(NUM_LEG * 6 + i) = (1.0 - estimated_contacts[i]) * (x(2) + fk_pos(2)) + estimated_contacts[i] * 0; // height z estimation
        }

        S = C * Pbar * C.transpose() + R;
        S = 0.5 * (S + S.transpose());

        error_y = y - yhat;
        Serror_y = S.fullPivHouseholderQr().solve(error_y);
        // auto solver = Eigen::FullPivHouseholderQR<Eigen::Matrix<T, MEAS_SIZE, MEAS_SIZE>>(S);
        // Serror_y = solver.solve(error_y);

        dx_forward = xbar - x;
        x = xbar + Pbar * C.transpose() * Serror_y;
        dx_backward = x - xbar;

        SC = S.fullPivHouseholderQr().solve(C);
        // SC = solver.solve(C);
        P = Pbar - Pbar * C.transpose() * SC * Pbar;
        P = 0.5 * (P + P.transpose());

        Serror_y = S.fullPivHouseholderQr().solve(Eigen::Matrix<T, MEAS_SIZE, 1>::Ones());

        // reduce position drift
        if (P.template block<3, 3>(0, 0).determinant() > 1e-6)
        {
            P.template block<3, 15>(0, 3).setZero();
            P.template block<15, 3>(3, 0).setZero();
            P.template block<3, 3>(0, 0) /= 10.0;
        }

        // final step
        // put estimated values back to A1CtrlStates& state
        // for (int i = 0; i < NUM_LEG; ++i)
        // {
        //     if (estimated_contacts[i] < 0.5)
        //     {
        //         state.estimated_contacts[i] = false;
        //     }
        //     else
        //     {
        //         state.estimated_contacts[i] = true;
        //     }
        // }
        //    std::cout << x.transpose() <<std::endl;
        // std::cout << x.head(6).transpose() << std::endl;
        state.body_lin_pos = x.template segment<3>(0);
        state.body_lin_vel = x.template segment<3>(3);
    }

    bool is_inited() { return filter_initialized; }

    void set_gravity(T g)
    {
        gravity = g;
    }

    Eigen::Matrix<T, STATE_SIZE, 1> getForward() {return dx_forward; }
    Eigen::Matrix<T, STATE_SIZE, 1> getBackward() {return dx_backward; }
    Eigen::Matrix<T, STATE_SIZE, 1> getState() {return x; }
    Eigen::Matrix<T, STATE_SIZE, 1> getStateNoise() {return Pbar.diagonal(); }
    Eigen::Matrix<T, MEAS_SIZE, 1> getMeasure() {return y; }
    Eigen::Matrix<T, MEAS_SIZE, 1> getPrediction() {return yhat; }
    Eigen::Matrix<T, MEAS_SIZE, 1> getError() {return error_y; }
    Eigen::Matrix<T, MEAS_SIZE, 1> getMeasureNoise() {return S.diagonal(); }
    Eigen::Matrix<T, MEAS_SIZE, 1> getUnitCorrection() {return Serror_y; }

    Eigen::Matrix<T, STATE_SIZE, 1> getQ() {return Q.diagonal(); }
    Eigen::Matrix<T, MEAS_SIZE, 1> getR() {return R.diagonal(); }
    Eigen::Matrix<T, STATE_SIZE, 1> getP() {return P.diagonal(); }
    Eigen::Matrix<T, MEAS_SIZE, 1> getS() {return S.diagonal(); }

private:
    bool filter_initialized = false;
    // state
    // 0 1 2 pos 3 4 5 vel 6 7 8 foot pos FL 9 10 11 foot pos FR 12 13 14 foot pos RL 15 16 17 foot pos RR
    Eigen::Matrix<T, STATE_SIZE, 1> x;             // estimation state
    Eigen::Matrix<T, STATE_SIZE, 1> xbar;          // estimation state after process update
    Eigen::Matrix<T, STATE_SIZE, STATE_SIZE> P;    // estimation state covariance
    Eigen::Matrix<T, STATE_SIZE, STATE_SIZE> Pbar; // estimation state covariance after process update
    Eigen::Matrix<T, STATE_SIZE, STATE_SIZE> A;    // estimation state transition
    Eigen::Matrix<T, STATE_SIZE, 3> B;             // estimation state transition
    Eigen::Matrix<T, STATE_SIZE, STATE_SIZE> Q;    // estimation state transition noise

    // observation
    // 0 1 2   FL pos residual
    // 3 4 5   FR pos residual
    // 6 7 8   RL pos residual
    // 9 10 11 RR pos residual
    // 12 13 14 vel residual from FL
    // 15 16 17 vel residual from FR
    // 18 19 20 vel residual from RL
    // 21 22 23 vel residual from RR
    // 24 25 26 27 foot height
    Eigen::Matrix<T, MEAS_SIZE, 1> y;           //  observation
    Eigen::Matrix<T, MEAS_SIZE, 1> yhat;        // estimated observation
    Eigen::Matrix<T, MEAS_SIZE, 1> error_y;     // estimated observation
    Eigen::Matrix<T, MEAS_SIZE, 1> Serror_y;    // S^-1*error_y
    Eigen::Matrix<T, MEAS_SIZE, STATE_SIZE> C;  // estimation state observation
    Eigen::Matrix<T, MEAS_SIZE, STATE_SIZE> SC; // S^-1*C
    Eigen::Matrix<T, MEAS_SIZE, MEAS_SIZE> R;   // estimation state observation noise
    // helper matrices
    Eigen::Matrix<T, 3, 3> eye3;               // 3x3 identity
    Eigen::Matrix<T, MEAS_SIZE, MEAS_SIZE> S;  // Innovation (or pre-fit residual) covariance
    Eigen::Matrix<T, STATE_SIZE, MEAS_SIZE> K; // kalman gain

    bool assume_flat_ground = false;

    // variables to process foot force
    T smooth_foot_force[4];
    T estimated_contacts[4];
    
    // params
    T process_noise_p_imu = 0.2;
    T process_noise_v_imu = 0.2;
    T process_noise_p_foot = 0.002;
    T sensor_noise_p_imu_rel_foot = 0.005;
    T sensor_noise_v_imu_rel_foot = 0.1;
    T sensor_noise_z_foot = 0.001;
    T high_suspect_num = 1000;

    // log
    Eigen::Matrix<T, STATE_SIZE, 1> dx_forward;
    Eigen::Matrix<T, STATE_SIZE, 1> dx_backward;

    T gravity = 9.75;
};
